import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Define the function and its derivative
def f(x):
    return x ** 2
optimizers = [
    tf.keras.optimizers.SGD(learning_rate=0.1),
    tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.96),
    tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.96, nesterov=True),
]
 
optimizer_names = ["Gradient Descent", "Gradient Descent with Momentum", "Gradient Descent with Nesterov Momentum"]
# Set up the plot
fig, axs = plt.subplots(1, 3, figsize=(15, 5))
x_vals = np.linspace(-10, 10, 100)
y_vals = f(x_vals)

# Initialize variables for each optimizer
tfx_list = [tf.Variable(-4.87) for _ in range(len(optimizers))]
opt_list = optimizers
visited_points = [[] for _ in range(len(optimizers))]

# Animation function
def update(frame, lines, scatters, titles):
    for i, opt in enumerate(opt_list):
        val = opt.minimize(lambda: f(tfx_list[i]), var_list=[tfx_list[i]])
        visited_points[i].append((tfx_list[i].numpy(), f(tfx_list[i]).numpy()))
        lines[i].set_data(x_vals, f(x_vals))
        scatters[i].set_offsets(visited_points[i])
        titles[i].set_text(f"{optimizer_names[i]}\nIteration: {frame+1}\nMinimum value:\n ({np.round(tfx_list[i].numpy(), 3)}, {np.round(f(tfx_list[i]).numpy(), 3)})")

# Set up the plots for each optimizer
lines_list = [ax.plot([], [])[0] for ax in axs]
scatters_list = [ax.scatter([], [], color='blue') for ax in axs]
titles_list = [ax.set_title("SGD") for ax in axs]

for ax in axs:
    ax.plot(x_vals, y_vals)

# Set up the animation
ani = FuncAnimation(fig, update, frames=100, fargs=(lines_list, scatters_list, titles_list), interval=100, repeat=False)

# Display the animation
from IPython.display import HTML
HTML(ani.to_jshtml())
WARNING:tensorflow:From C:\Users\temir\AppData\Local\Programs\Python\Python311\Lib\site-packages\keras\src\losses.py:2976: The name tf.losses.sparse_softmax_cross_entropy is deprecated. Please use tf.compat.v1.losses.sparse_softmax_cross_entropy instead.
_images/d396d2baaa01f3142052b9be600611b2a3f7bd38e8b276337882f9c80e4182f2.png